/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package com.icee.myth.bootstrap;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManagerFactory;
import javax.tools.*;
import javax.tools.JavaFileObject.Kind;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.security.KeyStore;
import java.security.SecureClassLoader;
import java.security.Security;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

/**
 *
 * @author liuxianke
 */
public class BootstrapClassLoader extends SecureClassLoader {
    //内部类
    private class BootstrapFileManager extends ForwardingJavaFileManager<JavaFileManager> {
        
        public BootstrapFileManager(JavaFileManager parentManager) {
            super(parentManager);
        }

        @Override
        public ClassLoader getClassLoader(Location location) {
            return BootstrapClassLoader.this;
        }

        @Override
        public Iterable<JavaFileObject> list(Location location,
                                             String packageName,
                                             Set<Kind> kinds,
                                             boolean recurse)
            throws IOException
        {
            List<JavaFileObject> retList = new ArrayList<JavaFileObject>();
            Iterator it = jclassFileObjectMap.entrySet().iterator();
            while (it.hasNext()) {
                Map.Entry<String, JavaFileObject> pair = (Map.Entry)it.next();
                String path = pair.getKey();
                if (path.substring(0, path.lastIndexOf(".")).equals(packageName)) {
                    retList.add(pair.getValue());
                }
            }
            
            if (retList.isEmpty()) {
                return super.list(location, packageName, kinds, recurse);
            } else {
                return retList;
            }
        }

        @Override
        public String inferBinaryName(Location location, JavaFileObject file) {
            String result;
            if (file instanceof BootstrapJavaFileObject) {
                result = file.getName().replaceFirst(".class", "").replaceAll("/", ".");
                if (result.startsWith(".")) {
                    result = result.replaceFirst(".", "");
                }
            } else {
                result = super.inferBinaryName(location, file);
            }
            return result;
        }
    }
    //内部类
    private class BootstrapJavaFileObject extends SimpleJavaFileObject {

        public final byte[] data;

        /**
        * Registers the compiled class object under URI
        * containing the class full name
        *
        * @param name
        *            Full name of the compiled class
        * @param kind
        *            Kind of the data. It will be CLASS in our case
        */
        public BootstrapJavaFileObject(String name, Kind kind, byte[] data) {
            super(URI.create("string:///" + name.replace('.', '/')
                + kind.extension), kind);

            this.data = data;
        }

        @Override
        public InputStream openInputStream() throws IOException {
            return new ByteArrayInputStream(data);
        }
    }

    private boolean initialized = false;
    private JavaFileManager fileManager = null;
    private HashMap<String, BootstrapJavaFileObject> jclassFileObjectMap = new HashMap<String, BootstrapJavaFileObject>();

    public BootstrapClassLoader(ClassLoader parent) {
        super(parent);
    }

    public JavaFileManager getFileManager() {
        if (fileManager == null) {
            fileManager = new BootstrapFileManager(ToolProvider.getSystemJavaCompiler().getStandardFileManager(null, null, null));
        }
        return fileManager;
    }

    private void getServerJAR() {
        // Get JAR from manager
        try {
            // 1. 建立与Manager Server的SSL连接
            String algorithm = Security.getProperty("ssl.TrustManagerFactory.algorithm");
            if (algorithm == null) {
                algorithm = "SunX509";
            }

            SSLContext clientContext = null;

            KeyStore ks = KeyStore.getInstance("JKS");
            ks.load(TrustStore.asInputStream(),
                    TrustStore.getTrustStorePassword());

            // Set up key manager factory to use our key store
            TrustManagerFactory tmf = TrustManagerFactory.getInstance(algorithm);
            tmf.init(ks);

            // Initialize the SSLContext to work with our key managers.
            clientContext = SSLContext.getInstance("TLS");
            clientContext.init(null, tmf.getTrustManagers(), null);

            String managerHost = System.getProperty("com.icee.myth.manager.host");
            String managerPort = System.getProperty("com.icee.myth.manager.port");

            SSLSocketFactory sslsocketfactory = clientContext.getSocketFactory();
            SSLSocket sslsocket = (SSLSocket) sslsocketfactory.createSocket(managerHost, Integer.valueOf(managerPort));

            // 2. 获取Manager Server传来的JAR并解开存入jclassDataMap中
            InputStream inputstream = sslsocket.getInputStream();
            JarInputStream jis = new JarInputStream(inputstream);

            JarEntry je = null;
            while ((je = jis.getNextJarEntry()) != null) {
                int entrySize = (int)je.getSize();
                if (entrySize>0) {
                    String jename = je.getName();
//                    System.out.println(jename);
                    String str = jename.replaceFirst(".class", "");
                    str = str.replaceAll("/", ".");

                    byte[] classbytes = new byte[entrySize];

                    int currentRead = 0;
                    while (currentRead<entrySize) {
                        currentRead += jis.read(classbytes, currentRead, entrySize-currentRead);
                    }

                    BootstrapJavaFileObject fileObject = new BootstrapJavaFileObject(str, Kind.CLASS, classbytes);
                    jclassFileObjectMap.put(str, fileObject);
                }
            }

            jis.close();

            inputstream.close();

            sslsocket.close();
        } catch (Exception exception) {
            exception.printStackTrace();
        }
    }

    @Override
    protected Class<?> findClass(String name)
            throws ClassNotFoundException {
        if (!initialized) {
            getServerJAR();
            initialized = true;
        }

        BootstrapJavaFileObject fileObject = jclassFileObjectMap.get(name);
        if (fileObject != null) {
            return super.defineClass(name, fileObject.data, 0, fileObject.data.length);
        } else {
            throw new ClassNotFoundException(name+" not found.");
        }
    }
}
